import rsa
import os
import hashlib

dir_path = os.path.dirname(os.path.realpath(__file__)).replace("\\", "/")

class antrsa:
    def __init__(self, key_size=2048, key_path=None, key_name=None):
        self.key_size = key_size
        self.key_name = key_name
        self.key_path = key_path
        if key_path is None:
            self.key_path = dir_path + "/keys/"
        if os.path.exists(self.key_path) == False:
            raise FileNotFoundError("Key does not exist")

    def generate_keys(self):
        (self.public_key, self.private_key) = rsa.newkeys(self.key_size)
        # Save the keys to files
        self.key_name = hashlib.md5(self.public_key.save_pkcs1().hex().encode()).hexdigest()
        if self.key_path is not None:
            rsa.save_key(self.public_key, f"{self.key_path}/{self.key_name}_public.pem")
            rsa.save_key(self.private_key, f"{self.key_path}/{self.key_name}_private.pem")
        else:
            if not os.path.exists(dir_path + "/keys"):
                os.mkdir(dir_path + "/keys")
            rsa.save_key(self.public_key, f"{dir_path}/{self.key_name}_public.pem")
            rsa.save_key(self.private_key, f"{dir_path}/{self.key_name}_private.pem")
        return self.public_key, self.private_key

    def encrypt(self, message):
        if self.public_key is None:
            raise ValueError("Public key not generated")
        return rsa.encrypt(message.encode(), self.public_key)

    def decrypt(self, encrypted_message):
        if self.private_key is None:
            raise ValueError("Private key not generated")